import os
from modules.dataset import DatasetModel
from modules.training import TrainingModel
from modules.collection import CollectionModel
from services.rtsp import RtspCollectionService
from modules.camera import CameraModel
from modules.image import ImageModel
import threading
import shutil
from datetime import datetime
from flask_jwt_extended import get_jwt_identity, jwt_required
from services.collect import CollectionService
from core import db, app
from common.resp import ajax_success, ajax_error, ajax_data, ajax_page

events = {}
def start_collection_thread(collection: CollectionModel, event: threading.Event):
    with app.app_context():
        
        #开始
        CollectionModel.query.filter_by(id=collection.id).update({
            'status': '1',
            'start_time': datetime.now()
        })
        db.session.commit()

        camera = CameraModel.query.get_or_404(collection.camera_id)
        if camera.type == 'rtsp' or camera.type == 'playback':
            service = RtspCollectionService(collection, event)
            service.start();

        elif camera.type == 'local':
            service = CollectionService(collection)
            service.start();
        else:
            logger.error(f"Camera type {camera.type} not supported.")
    
        #开始
        CollectionModel.query.filter_by(id=collection.id).update({
            'status': '4',
            'end_time': datetime.now()
        })
        db.session.commit()
        
@app.route("/collections/start/<int:id>", methods=['GET'])
@jwt_required()
def start_collection(id):
    account = get_jwt_identity();
    collection = CollectionModel.query.get_or_404(id)
    if not collection:
        return ajax_error('采集任务不存在')
    
    event = threading.Event()
    events[collection.id] = event    
    thread = threading.Thread(target=start_collection_thread, args=(collection,event,))
    thread.start()
    return ajax_success('采集任务启动成功')

@app.route("/collections/stop/<int:id>", methods=['GET'])
@jwt_required()
def stop_collection(id):
    account = get_jwt_identity();

    collection = CollectionModel.query.get_or_404(id)
    if not collection:
        return ajax_error('采集任务不存在')

    if collection.status != '1':
        return ajax_error('采集任务不是运行状态')
    
    event = events.get(collection.id)
    if event:
        event.set()
        
    collection.status = '4'
    collection.end_time = datetime.now()
    db.session.add(collection)
    db.session.commit()
    return ajax_success('采集任务停止成功')

@app.route("/datasets/write/<int:id>", methods=['GET'])
@jwt_required()
def write_dataset_image(id):
    dataset = DatasetModel.query.get_or_404(id)
    if not dataset:
        return ajax_error('数据集不存在')

    dsPath = os.path.join(app.config['DATASET_PATH'], dataset.path)
    if not os.path.exists(dsPath):
        os.makedirs(dsPath)

    filePath = os.path.join(dsPath, f"yolo.yaml")
    with open(filePath, 'w') as f:
        nc = len(dataset.classes.split('\n'))
        f.write(f"nc: {nc}\n")
        f.write(f"names:\n {dataset.classes}\n")

        f.write(f"path: {dsPath}\n")
        trainPath = os.path.join(dsPath, 'images/train')
        if not os.path.exists(trainPath):
            os.makedirs(trainPath)
            os.makedirs(os.path.join(dsPath, 'labels/train'))
        f.write(f"train: {os.path.abspath(os.path.join(dsPath, 'images/train'))}\n")    

        valPath = os.path.join(dsPath, 'images/val')
        if not os.path.exists(valPath):
            os.makedirs(valPath)
            os.makedirs(os.path.join(dsPath, 'labels/val'))
        f.write(f"val: {os.path.abspath(os.path.join(dsPath, 'images/val'))}\n")

        testPath = os.path.join(dsPath, 'images/test')
        if not os.path.exists(testPath):
            os.makedirs(testPath)
            os.makedirs(os.path.join(dsPath, 'labels/test'))
        f.write(f"test: {os.path.abspath(os.path.join(dsPath, 'images/test'))}\n")
        f.flush()

    images = ImageModel.query.filter(ImageModel.ds_id == dataset.id).all();
    for image in images:
        imgPath = os.path.join(dsPath, image.name)
        if not os.path.exists(imgPath):
            continue
        
        print("move", imgPath, "to", os.path.join(dsPath, image.category, image.name))
        shutil.move(os.path.abspath(imgPath), os.path.abspath(os.path.join(dsPath, 'images',image.category, image.name)))
        labelPath = os.path.join(dsPath, f"labels/{image.category}/{image.name.split('.')[0]}.txt")
        with open(labelPath, 'w') as f:
            for annotation in image.annotations:
                # 转换为YOLO格式（归一化的中心点坐标和宽高）
                x_center = annotation['x'] + annotation['width'] / 2 
                y_center = annotation['y'] + annotation['height'] / 2 
                bbox_width = annotation['width']
                bbox_height = annotation['height']

                # 写入文件
                f.write(f"{annotation['name']} {x_center:.6f} {y_center:.6f} {bbox_width:.6f} {bbox_height:.6f}\n")
            f.flush()

    dataset.status = '4'
    db.session.add(dataset)
    db.session.commit()
    return ajax_success('数据集分类写入成功')


def start_train_thread(training: TrainingModel):
    with app.app_context():
        camera = CameraModel.query.get_or_404(training.camera_id)
        if camera.type == 'rtsp' or camera.type == 'playback':
            service = RtspCollectionService(training)
            service.start();
        elif camera.type == 'local':
            service = CollectionService(training)
            service.start();
        else:
            logger.error(f"Camera type {camera.type} not supported.")

@app.route("/training/start/<int:id>", methods=['GET'])
@jwt_required()
def start_training(id):
    account = get_jwt_identity();
    training = TrainingModel.query.get_or_404(id)
    if not training:
        return ajax_error('训练任务不存在')

    thread = threading.Thread(target=start_train_thread, args=(training,))
    thread.start()
    return ajax_success('训练任务启动成功')

@app.route("/training/stop/<int:id>", methods=['GET'])
@jwt_required()
def stop_training(id):
    account = get_jwt_identity();

    training = TrainingModel.query.get_or_404(id)
    if not training:
        return ajax_error('训练任务不存在')

    training.status = '4'
    training.end_time = datetime.now()
    db.session.add(training)
    db.session.commit()
    return ajax_success('训练任务停止成功')
